package com.darrenchan.leetcode;

import java.util.*;

/**
 * @Desc
 * @Author chenchi03
 * @CreateTime 2020-03-24 22:15
 */
public class TopK {

    Map<String, Integer> map;
    TreeSet<String> topk;
    int k;
    /*
     * @param k: An integer
     */
    public TopK(int k) {
        map = new HashMap<>();
        topk = new TreeSet<>((o1, o2) -> {
            if(o1.equals(o2)) {
                return 0;
            }
            if(!map.get(o1).equals(map.get(o2))) {
                return map.get(o1) - map.get(o2);
            }
            return o2.compareTo(o1);
        });

        this.k = k;
    }

    /*
     * @param word: A string
     * @return: nothing
     */
    public void add(String word) {
        if(!map.containsKey(word)) {
            map.put(word, 1);
        } else {
            if(topk.contains(word)) {
                topk.remove(word);
            }
            map.put(word, map.get(word) + 1);
        }

        topk.add(word);
        if(topk.size() > k) {
            topk.pollFirst();
        }
    }

    /*
     * @return: the current top k frequent words.
     */
    public List<String> topk() {
        List<String> res = new ArrayList<>();
        for (String re : topk) {
            res.add(0, re);
        }
        return res;
    }
}
